import torch
from torch import nn
from torch.nn import functional as F
# https://zh-v2.d2l.ai/chapter_convolutional-modern/resnet.html#id4

class Residual(nn.Module):  #@save
    def __init__(self, input_channels, num_channels,
                 use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, num_channels,
                               kernel_size=3, padding=1, stride=strides)
        self.conv2 = nn.Conv2d(num_channels, num_channels,
                               kernel_size=3, padding=1)
        # 当 use_1x1conv = True 时，添加通过1*1卷积调整通道和分辨率。
        # 当 use_1x1conv = False 时，应用ReLU非线性函数之前，将输入添加到输出。
        if use_1x1conv:
            self.conv3 = nn.Conv2d(input_channels, num_channels,
                                   kernel_size=1, stride=strides)
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)

    def forward(self, X):
        # 每个残差块有2个卷积层（不包括恒等映射的1*1卷积层）
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)


class ResNet18(nn.Module):
    def __init__(self, residual_block = Residual, num_classes=1000):
        super(ResNet18, self).__init__()
        self.b1 = nn.Sequential(nn.Conv2d(in_channels = 3, out_channels = 64, 
                                    kernel_size=7, stride=2, padding=3),
                        nn.BatchNorm2d(64), nn.ReLU(),
                        nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        # 每个残差模块使用2个残差块，所以每个残差模块有4个卷积层
        # ResNet18: 使用4个由残差块组成的模块，每个模块使用若干个同样输出通道数的残差块。加上第一个7*7卷积层和最后一个全连接层，共有18层。 
        # 第一个模块的通道数同输入通道数一致。 由于之前已经使用了步幅为2的最大汇聚层，所以无须减小高和宽。
        # 之后的每个模块在第一个残差块里将上一个模块的通道数翻倍，并将高和宽减半。
        # 分辨率降低，通道数量增加，直到全局平均汇聚层聚集所有特征。
        self.b25 = nn.Sequential(*self.resnet_block(residual_block, input_channels = 64, num_channels = 64, num_residuals = 2, first_block=True))
        self.b69 = nn.Sequential(*self.resnet_block(residual_block, input_channels = 64, num_channels = 128, num_residuals = 2))
        self.b1013 = nn.Sequential(*self.resnet_block(residual_block, input_channels = 128, num_channels = 256, num_residuals = 2))
        self.b1417 = nn.Sequential(*self.resnet_block(residual_block, input_channels = 256, num_channels = 512, num_residuals = 2))
        self.b18 = nn.Sequential(nn.AdaptiveAvgPool2d((1,1)),
                            nn.Flatten(), nn.Linear(512, num_classes))   

    def resnet_block(self, residual_block, input_channels, num_channels, num_residuals, first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
            # 非第1个残差模块里面的第1个残差块用1*1卷积层
                blk.append(residual_block(input_channels, num_channels,
                                    use_1x1conv=True, strides=2))
            else: 
            # 第1个残差模块里面的2个残差块都不用1*1卷积层
            # 非第1个残差模块里面的第2个残差块用1*1卷积层
                blk.append(residual_block(num_channels, num_channels))
        return blk

    def forward(self, x):
        x = self.b1(x)
        x = self.b25(x)
        x = self.b69(x)
        x = self.b1013(x)
        x = self.b1417(x)
        x = self.b18(x)
        return x
        
    
resnet18 = ResNet18(residual_block = Residual, num_classes=7)
net = nn.Sequential(resnet18.b1, resnet18.b25, resnet18.b69, resnet18.b1013, resnet18.b1417, resnet18.b18)    
# 输入的大小=(批量大小, 通道数, 行数, 列数)
X = torch.rand(size=(31, 3, 224, 224))
# 打印网络的大小
# 步长为2的话，高和宽减半
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__,'output shape: ', X.shape) 









# def resnet_block(input_channels, num_channels, num_residuals, first_block=False):
#     blk = []
#     for i in range(num_residuals):
#         if i == 0 and not first_block:
#         # 非第1个残差模块里面的第1个残差块用1*1卷积层
#             blk.append(Residual(input_channels, num_channels,
#                                 use_1x1conv=True, strides=2))
#         else: 
#         # 第1个残差模块里面的2个残差块都不用1*1卷积层
#         # 非第1个残差模块里面的第2个残差块用1*1卷积层
#             blk.append(Residual(num_channels, num_channels))
#     return blk
    
# b1 = nn.Sequential(nn.Conv2d(in_channels = 3, out_channels = 64, 
#                              kernel_size=7, stride=2, padding=3),
#                    nn.BatchNorm2d(64), nn.ReLU(),
#                    nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
# # 每个残差模块使用2个残差块，所以每个残差模块有4个卷积层
# # ResNet18: 加上第一个7*7卷积层和最后一个全连接层，共有18层。 
# b2 = nn.Sequential(*resnet_block(input_channels = 64, num_channels = 64, num_residuals = 2, first_block=True))
# b3 = nn.Sequential(*resnet_block(input_channels = 64, num_channels = 128, num_residuals = 2))
# b4 = nn.Sequential(*resnet_block(input_channels = 128, num_channels = 256, num_residuals = 2))
# b5 = nn.Sequential(*resnet_block(input_channels = 256, num_channels = 512, num_residuals = 2))
# net = nn.Sequential(b1, b2, b3, b4, b5,
#                     nn.AdaptiveAvgPool2d((1,1)),
#                     nn.Flatten(), nn.Linear(512, 10))

# X = torch.rand(size=(256, 3, 224, 224))
# for layer in net:
#     X = layer(X)
#     print(layer.__class__.__name__,'output shape:\t', X.shape)